import numpy as np
from scipy.fft import fft
import pyaudio
import time
from matplotlib import font_manager
import serial
import serial.tools.list_ports
from tkinter import *
from tkinter import ttk
import threading
import queue

count = 0
start_time = None

class AudioSpectrumAnalyzer:
    def __init__(self):
        self.serial_port = None
        self.is_running = False  # 初始状态应该是False
        self.audio_queue = queue.Queue()  # 原始音频数据队列
        
        # 频段范围（Hz）- 与实际FFT bin分配保持一致
        self.freq_bands = {
            'low': (0, 188),      # 低频：0-187.5Hz (bin 0-7)
            'mid': (211, 563),    # 中频：210.9-562.5Hz (bin 9-24)  
            'high': (609, 1969)   # 高频：609.4-1968.8Hz (bin 26-84)
        }
        
        # 各频段的值
        self.band_values = {
            'low': 0,
            'mid': 0,
            'high': 0
        }
        
        self.setup_gui()
        
    def setup_gui(self):
        """设置GUI界面"""
        self.root = Tk()
        self.root.title("灵动狗眼睛控制")
        self.root.geometry("400x500")
        
        # 设置串口选择框
        frame = ttk.Frame(self.root, padding="10")
        frame.grid(row=0, column=0, sticky=(W, E, N, S))
        
        ttk.Label(frame, text="选择串口:").grid(row=0, column=0, sticky=W)
        self.port_var = StringVar()
        self.port_combo = ttk.Combobox(frame, textvariable=self.port_var)
        self.port_combo.grid(row=0, column=1, sticky=(W, E))
        
        # 波特率选择
        ttk.Label(frame, text="波特率:").grid(row=1, column=0, sticky=W)
        self.baud_var = StringVar(value="115200")
        baud_rates = ["9600", "19200", "38400", "57600", "115200"]
        self.baud_combo = ttk.Combobox(frame, textvariable=self.baud_var, values=baud_rates)
        self.baud_combo.grid(row=1, column=1, sticky=(W, E))
        
        # 连接按钮
        self.connect_button = ttk.Button(frame, text="连接", command=self.toggle_connection)
        self.connect_button.grid(row=2, column=0, columnspan=2, pady=10)
        
        # 创建Canvas用于显示频谱
        self.canvas = Canvas(self.root, width=300, height=300, bg='black')
        self.canvas.grid(row=1, column=0, padx=50, pady=20)
        
        # 创建三个频段的柱状图
        bar_width = 50
        spacing = 25
        self.spectrum_bars = {
            'low': self.canvas.create_rectangle(50, 300, 50+bar_width, 300, fill='#FF0000'),
            'mid': self.canvas.create_rectangle(125, 300, 125+bar_width, 300, fill='#00FF00'),
            'high': self.canvas.create_rectangle(200, 300, 200+bar_width, 300, fill='#0000FF')
        }
        
        # 创建频段标签
        self.canvas.create_text(75, 310, text="低频", fill='white')
        self.canvas.create_text(150, 310, text="中频", fill='white')
        self.canvas.create_text(225, 310, text="高频", fill='white')
        
        # 创建刻度
        for i in range(11):
            y = 300 - i * 30
            self.canvas.create_line(45, y, 255, y, fill='gray')
            self.canvas.create_text(35, y, text=str(i*25), fill='white')
        
        # 创建按钮样式
        self.style = ttk.Style()
        self.style.configure("Red.TButton", foreground="red")
        
        # 开始/停止按钮
        self.start_button = ttk.Button(self.root, text="开始监听", command=self.toggle_monitor)
        self.start_button.grid(row=2, column=0, pady=10)
        
        # 刷新串口列表
        self.refresh_ports()
        
        # 定期刷新串口列表
        self.root.after(1000, self.refresh_ports)
        
    # 刷新可用串口列表
    def refresh_ports(self):
        """刷新可用串口列表"""
        ports = [port.device for port in serial.tools.list_ports.comports()]
        self.port_combo['values'] = ports
        if ports and not self.port_var.get():
            self.port_var.set(ports[0])
        self.root.after(1000, self.refresh_ports)

    # 更新频谱显示
    def update_spectrum_display(self):
        """更新频谱显示"""
        if self.is_running:
            # 更新三个频段的柱状图
            for band, bar in self.spectrum_bars.items():
                value = self.band_values[band]
                height = min(300, max(0, value * 1.2))  # 值映射到高度
                y1 = 300 - height
                x1, _, x2, _ = self.canvas.coords(bar)
                self.canvas.coords(bar, x1, y1, x2, 300)
            
            # 每33ms更新一次（约30fps）
            self.root.after(33, self.update_spectrum_display)

    # 音频处理线程函数
    def process_audio(self):
        """音频处理线程函数"""
        SAMPLE_RATE = 48000
        WINDOW_SIZE = 2048
        
        # 计算频率分辨率
        freq_resolution = SAMPLE_RATE / WINDOW_SIZE
        print(f"频率分辨率: {freq_resolution:.2f} Hz/bin")
        
        # 从freq_bands计算FFT bin范围，确保一致性
        freq_bins = {}
        for band, (low, high) in self.freq_bands.items():
            low_bin = int(low / freq_resolution)
            high_bin = int(high / freq_resolution)
            freq_bins[band] = (low_bin, high_bin)
            actual_low = low_bin * freq_resolution
            actual_high = high_bin * freq_resolution
            print(f"{band}频段: {low}-{high}Hz -> bin {low_bin}-{high_bin} (实际: {actual_low:.1f}-{actual_high:.1f}Hz)")
        
        # 初始化数据
        audio_data = np.zeros(WINDOW_SIZE, dtype=np.float32)
        window = np.hanning(WINDOW_SIZE)
        
        # 噪声门限和动态范围参数
        noise_threshold_db = -60  # 噪声门限 (dB)
        min_db = -60              # 映射的最小dB值
        max_db = -15              # 映射的最大dB值
        
        while self.is_running:
            try:
                # 从队列获取音频数据
                try:
                    while not self.audio_queue.empty():
                        data = np.frombuffer(self.audio_queue.get_nowait(), dtype=np.float32)
                        audio_data = np.roll(audio_data, -len(data))
                        audio_data[-len(data):] = data
                except queue.Empty:
                    pass

                # 应用窗函数
                windowed_data = audio_data * window
                
                # 计算FFT并进行归一化
                fft_data_raw = np.abs(fft(windowed_data))[:WINDOW_SIZE//2]
                fft_data = fft_data_raw / (WINDOW_SIZE / 2) # <-- 关键的归一化步骤

                # 找到最强的频率分量（用于调试）
                max_bin = np.argmax(fft_data)
                max_freq = max_bin * freq_resolution
                
                # 计算各频段的能量
                for band, (start_bin, end_bin) in freq_bins.items():
                    # 提取频段数据
                    band_data = fft_data[start_bin:end_bin]
                    if len(band_data) == 0:
                        self.band_values[band] = 0
                        continue
                        
                    # 计算均方根值
                    rms = np.sqrt(np.mean(np.square(band_data)))
                    
                    # 转换为分贝值
                    if rms > 1e-9: # 避免log(0)
                        db = 20 * np.log10(rms)
                        
                        # 应用噪声门限
                        if db < noise_threshold_db:
                            value = 0
                        else:
                            # 映射到0-1范围
                            normalized = (db - min_db) / (max_db - min_db)
                            value = max(0, min(1, normalized)) # 限制在0-1
                        
                        self.band_values[band] = int(value * 255)
                    else:
                        self.band_values[band] = 0
                
                time.sleep(0.01) # 稍微增加休眠时间，降低CPU占用
                    
            except Exception as e:
                print(f"处理音频数据错误: {str(e)}")
                time.sleep(0.01)

    # 串口发送线程函数
    def send_spectrum_data(self):
        """串口发送线程函数"""
        global count, start_time
        last_send_time = time.perf_counter()
        send_interval = 0.033  # 33ms发送间隔
        
        while self.is_running:
            try:
                if self.serial_port and self.serial_port.is_open:
                    current_time = time.perf_counter()
                    if current_time - last_send_time >= send_interval:
                        # 获取频段值
                        low = self.band_values['low']
                        mid = self.band_values['mid']
                        high = self.band_values['high']
                        
                        level = 0x64    # 眼睛亮度 0-100
                        data = bytes([
                            0xbe,  # 起始字节
                            0x11,  # 固定字节
                            0x0a,  # 固定字节
                            0xf1,  # 固定字节
                            level, # 眼睛亮度 0-100
                            low,   # 低频值 (0-255)，眼睛红色
                            mid,   # 中频值 (0-255)，眼睛绿色
                            high,  # 高频值 (0-255)，眼睛蓝色
                            0xeb,  # 固定字节
                            0xed   # 固定字节
                        ])
                        
                        if start_time is None:
                            start_time = time.perf_counter()
                        
                        self.serial_port.write(data)
                        count += 1
                        
                        # 每秒打印一次统计信息
                        elapsed_time = time.perf_counter() - start_time
                        if count % 20 == 0:  # 每20次打印一次
                            print(f"L:{low:02X} M:{mid:02X} H:{high:02X}")  # 以十六进制格式打印
                            # 打印实际发送频率
                            actual_rate = count / elapsed_time
                            print(f"实际发送频率: {actual_rate:.1f}Hz")
                        
                        last_send_time = current_time
                
                time.sleep(0.001)  # 短暂休眠，避免CPU占用过高
                    
            except Exception as e:
                print(f"发送数据失败: {str(e)}")
                time.sleep(0.1)  # 发生错误时稍微多等待一下

    # 切换串口连接状态
    def toggle_connection(self):
        """切换串口连接状态"""
        if self.serial_port is None:
            try:
                port = self.port_var.get()
                baud = int(self.baud_var.get())
                self.serial_port = serial.Serial(port, baud, timeout=1)
                self.connect_button.configure(text="断开")
                print(f"已连接到 {port}, 波特率: {baud}")
            except Exception as e:
                print(f"连接失败: {str(e)}")
        else:
            try:
                self.serial_port.close()
                self.serial_port = None
                self.connect_button.configure(text="连接")
                print("已断开连接")
            except Exception as e:
                print(f"断开连接失败: {str(e)}")

    def toggle_monitor(self):
        """切换监听状态"""
        if not self.is_running:
            self.start_monitor()
        else:
            self.stop_monitor()
    
    def start_monitor(self):
        """启动音频监测"""
        self.is_running = True
        
        # 启动音频捕获线程
        audio_thread = threading.Thread(target=self.capture_audio)
        audio_thread.daemon = True
        audio_thread.start()
        
        # 启动音频处理线程
        process_thread = threading.Thread(target=self.process_audio)
        process_thread.daemon = True
        process_thread.start()
        
        # 启动串口发送线程
        serial_thread = threading.Thread(target=self.send_spectrum_data)
        serial_thread.daemon = True
        serial_thread.start()
        
        # 启动显示更新
        self.update_spectrum_display()
        
        # 更新按钮状态
        self.start_button.configure(text="停止监听", style="Red.TButton")
        
        # 延迟检查音频捕获是否成功
        self.root.after(500, self.check_audio_status)
        
    def check_audio_status(self):
        """检查音频捕获状态，如果失败则恢复按钮状态"""
        if not self.is_running:
            # 音频捕获失败，恢复按钮状态
            self.start_button.configure(text="开始监听", style="TButton")
            print("音频捕获失败，请检查音频设备设置")
        
    def stop_monitor(self):
        """停止音频监测"""
        self.is_running = False
        
        # 重置频段值
        for band in self.band_values:
            self.band_values[band] = 0
            
        # 更新按钮状态
        self.start_button.configure(text="开始监听")
        self.start_button.configure(style="TButton")  # 恢复默认样式

    def capture_audio(self):
        """音频捕获线程"""
        SAMPLE_RATE = 48000
        WINDOW_SIZE = 2048
        FORMAT = pyaudio.paFloat32

        try:
            audio = pyaudio.PyAudio()
            # 查找立体声混音设备
            stereo_mix_index = None
            for i in range(audio.get_device_count()):
                dev_info = audio.get_device_info_by_index(i)
                if "立体声混音" in dev_info['name'] and dev_info['maxInputChannels'] > 0:
                    stereo_mix_index = i
                    break

            if stereo_mix_index is None:
                print("\n错误：未找到立体声混音设备！")
                self.is_running = False  # 设置为False，让其他线程也停止
                return

            def audio_callback(in_data, frame_count, time_info, status):
                if self.is_running:
                    self.audio_queue.put(in_data)
                return (in_data, pyaudio.paContinue)

            stream = audio.open(
                format=FORMAT,
                channels=2,
                rate=SAMPLE_RATE,
                input=True,
                input_device_index=stereo_mix_index,
                frames_per_buffer=WINDOW_SIZE//2,
                stream_callback=audio_callback
            )
            
            stream.start_stream()
            
            while self.is_running:
                time.sleep(0.1)
            
            stream.stop_stream()
            stream.close()
            audio.terminate()
            
        except Exception as e:
            print(f"音频捕获错误: {str(e)}")
            self.is_running = False  # 出错时也要设置为False

def main():
    app = AudioSpectrumAnalyzer()
    app.root.mainloop()

if __name__ == "__main__":
    main() 